/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once

#include <baja/graphics/dx11.1/details/common.hpp>
#include <baja/graphics/dx11.1/graphics_fwd.hpp>
#include <baja/graphics/dx11.1/details/enum_mappings.hpp>
#include <baja/graphics/dx11.1/details/util.hpp>

namespace baja { namespace graphics { namespace dx11_1 { namespace details {

// wraps a 2D texture
class texture2d : public b::iunknown_provider
{
public:
    texture2d() : _locked(false) {}
    static std::shared_ptr<texture2d> create(ID3D11Texture2D* t);
    static std::shared_ptr<texture2d> create(ID3D11View *view, const texture_type type);
    static std::shared_ptr<texture2d> create(const std::vector<WRL::ComPtr<ID3D11View>>& views);    

    // b::iunknown_provider
    size_t getInterfaceCount();
    WRL::ComPtr<IUnknown> getInterface(size_t index);

    uint32 getWidth();
    uint32 getHeight();
    pixel_format getPixelFormat();
    uint32 getArraySize();
    uint32  getMipLevelCount();

    size_t calculateSize(
        const uint32 sourceWidth,
        const uint32 sourceHeight
        );

    void copyToSystemMemory(
        const uint32 mipLevel,
        const uint32 sourceX,
        const uint32 sourceY,
        const uint32 sourceWidth,
        const uint32 sourceHeight,
        memory* p
        );

    void 
    copyFromSystemMemory(
        const uint32 mipLevel,
        const uint32 destX,
        const uint32 destY,
        const uint32 dataX,
        const uint32 dataY,
        const uint32 dataWidth,
        const uint32 dataHeight,
        const memory* data,
        const uint32 dataSizeInBytes,
        const uint32 dataBytesPerPixel,
        const uint32 pitch
        );
    
    uint32 getBitsPerPixel();

    void lock( 
        const uint32 level,
        const lock_type lockType,
        memory** data,
        uint32* pitch
        ); 
        
    void unlock();

private:

    void initialize(ID3D11Texture2D* texture2d);
    void initialize(ID3D11View* view);    
    void initialize(const std::vector<WRL::ComPtr<ID3D11View>>& views);

private:

    std::vector<WRL::ComPtr<ID3D11View>> _views;
    WRL::ComPtr<ID3D11Texture2D> _texture;
    WRL::ComPtr<ID3D11Texture2D> _surrogate;

    uint32 _lockedLevel;
    bool _lockUsingSurrogate;
    bool _locked;
};

inline std::shared_ptr<texture2d> texture2d::create(ID3D11Texture2D* t)
{
    std::shared_ptr<texture2d> returnValue = std::make_shared<texture2d>();
    returnValue->initialize(t);
    return returnValue;
}

inline std::shared_ptr<texture2d> texture2d::create(ID3D11View *view, const texture_type type)
{
    std::shared_ptr<texture2d> returnValue = std::make_shared<texture2d>();
    std::vector<WRL::ComPtr<ID3D11View>> views;
    views.resize(4);
    views[static_cast<int>(type)] = view;
    returnValue->initialize(views);
    return returnValue;
}

inline std::shared_ptr<texture2d> texture2d::create(const std::vector<WRL::ComPtr<ID3D11View>>& views)
{
    std::shared_ptr<texture2d> returnValue = std::make_shared<texture2d>();
    returnValue->initialize(views);
    return returnValue;
}
    
inline size_t texture2d::getInterfaceCount()
{
    return _views.size() + 1; // index 0 is texture2d interface
}
    
inline WRL::ComPtr<IUnknown> texture2d::getInterface(size_t index)
{
    if (index == 0)
    {
        WRL::ComPtr<IUnknown> p;
        win32_result(_texture.As(&p));
        return p;
    }
    else
    {
        BAJA_CHECK_ARG(index <= _views.size());
        WRL::ComPtr<IUnknown> p;
        win32_result((HRESULT)_views[index - 1].As(&p));
        return p;
    }
}
    
// gets width of texture
inline uint32 texture2d::getWidth() 
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);

    return desc.Width;
}
    
// gets height of texture
inline uint32 texture2d::getHeight()
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);

    return desc.Height;
}
    
// gets pixel format of texture
inline pixel_format texture2d::getPixelFormat() 
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);

    return fromDXGI_FORMAT(desc.Format);
}

// gets texture array size
inline uint32 texture2d::getArraySize()
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);
         
    return desc.ArraySize;
}
    
// gets mip level count
inline uint32 texture2d::getMipLevelCount()
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);
         
    return desc.MipLevels;
}

inline size_t texture2d::calculateSize(
    const uint32 sourceWidth,
    const uint32 sourceHeight
    )
{
    uint32 bytesPerPixel = this->getBitsPerPixel() / 8;

    return bytesPerPixel * sourceWidth * sourceHeight;
}

inline void
texture2d::copyToSystemMemory(
    const uint32 mipLevel,
    const uint32 sourceX,
    const uint32 sourceY,
    const uint32 sourceWidth,
    const uint32 sourceHeight,
    memory* p
    )
{
    // get texture description
    D3D11_TEXTURE2D_DESC textureDesc;
    _texture->GetDesc(&textureDesc);

    // get device
    WRL::ComPtr<ID3D11Device> device;
    _texture->GetDevice(device.GetAddressOf());

    // get device context
    WRL::ComPtr<ID3D11DeviceContext> deviceContext;
    device->GetImmediateContext(deviceContext.GetAddressOf());

    uint32 levelWidth = textureDesc.Width >> mipLevel;
    if (levelWidth == 0)
    {
        levelWidth = 1;
    }

    uint32 levelHeight = textureDesc.Height >> mipLevel;
    if (levelHeight == 0)
    {
        levelHeight = 1;
    }

    win32_result wr;

    // data validation
    if ((sourceX + sourceWidth) > levelWidth || 
        (sourceY + sourceHeight) > levelHeight)
    {
        wr = E_INVALIDARG;
    }

    WRL::ComPtr<ID3D11Texture2D> textureToLock;
    uint32 levelToLock = 0;
    bool copyWholeTexture = false;

    bool canReadLock = 
        ((textureDesc.CPUAccessFlags & D3D11_CPU_ACCESS_READ) == D3D11_CPU_ACCESS_READ) &&
        !_locked &&
        (mipLevel < textureDesc.MipLevels);

    if (canReadLock)
    {
        textureToLock = _texture;
        levelToLock = mipLevel;

        if (sourceX == 0 && sourceY == 0 && sourceWidth == levelWidth && sourceHeight == levelHeight)
        {
            copyWholeTexture = true;
        }
        else
        {
            copyWholeTexture = false;
        }
    }
    else
    {
        // Only going to create a surrogate texture for the level we're locking, 
        // so the texture we end up copying from just has 1 mip level.
                
        levelToLock = 0;
        copyWholeTexture = true;

        D3D11_TEXTURE2D_DESC surrogateDesc;
        surrogateDesc.Width = sourceWidth;
        surrogateDesc.Height = sourceHeight;
        surrogateDesc.MipLevels = 1;
        surrogateDesc.ArraySize = 1;
        surrogateDesc.Format = textureDesc.Format;
        surrogateDesc.SampleDesc.Count = textureDesc.SampleDesc.Count;
        surrogateDesc.SampleDesc.Quality = textureDesc.SampleDesc.Quality;
        surrogateDesc.Usage = D3D11_USAGE_STAGING;
        surrogateDesc.BindFlags = 0;
        surrogateDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ|D3D11_CPU_ACCESS_WRITE;
        surrogateDesc.MiscFlags = 0;

        wr = device->CreateTexture2D(&surrogateDesc, NULL, textureToLock.GetAddressOf());
        D3D11_BOX sourceRegion;
        sourceRegion.left = sourceX;
        sourceRegion.right = sourceX + sourceWidth;
        sourceRegion.top = sourceY;
        sourceRegion.bottom = sourceY + sourceHeight;
        sourceRegion.front = 0;
        sourceRegion.back = 1;

        deviceContext->CopySubresourceRegion(
            textureToLock.Get(),
            levelToLock,
            0,
            0,
            0,
            _texture.Get(),
            mipLevel,
            &sourceRegion
            );
    }
        
    bool isLocked = false;
    D3D11_MAPPED_SUBRESOURCE mappedResource;
    
    withFinally::run(
    [&]()
    {
        wr = deviceContext->Map(
            textureToLock.Get(),
            levelToLock,
            D3D11_MAP_READ,
            0,
            &mappedResource
            );
            
        uint32 bytesPerPixel = 0;
        uint32 bitsPerPixel = 0;

        bitsPerPixel = this->getBitsPerPixel();
        bytesPerPixel = bitsPerPixel / 8;

        isLocked = true;
                
        uint32 pitch = bytesPerPixel * sourceWidth;
        uint32 dataSizeInBytes = pitch * sourceHeight;

        if (mappedResource.RowPitch < pitch)
        {
            wr = E_FAIL;
        }
        else if (mappedResource.RowPitch > pitch)
        {
            // D3D11 padded the texture, so we need to copy row by row 
            // to get the actual data.
            copyWholeTexture = false;
        }
            
        if (copyWholeTexture)
        {
            memcpy(p, mappedResource.pData, dataSizeInBytes);
        }
        else
        {
            // Copy row by row to get the specified rectangle from the texture.

            uint32 destBytesPerRow = pitch;
            memory* copyToPos = p;
            for (uint32 j = sourceY; j < sourceHeight; j++)
            {
                // NOTE: Used mappedResource.RowPitch since we are indexing into the texture,
                // and the data may be padded.
                uint32 textureIndex = (j * mappedResource.RowPitch) + (sourceX * bytesPerPixel);
                memcpy(copyToPos, (memory*)mappedResource.pData + textureIndex, destBytesPerRow);
                copyToPos += destBytesPerRow;
            }
        }
    },
    [&]()
    {
        if (isLocked)
        {
            deviceContext->Unmap(textureToLock.Get(), levelToLock);
        }
    });
}

inline void 
texture2d::copyFromSystemMemory(
    const uint32 mipLevel,
    const uint32 destX,
    const uint32 destY,
    const uint32 dataX,
    const uint32 dataY,
    const uint32 dataWidth,
    const uint32 dataHeight,
    const memory* data,
    const uint32 dataSizeInBytes,
    const uint32 dataBytesPerPixel,
    const uint32 pitch
    )
{
    BAJA_CHECK_ARG(data != nullptr);
        
    // get texture description
    D3D11_TEXTURE2D_DESC textureDesc;
    _texture->GetDesc(&textureDesc);
        
    WRL::ComPtr<ID3D11Device> device;
    _texture->GetDevice(device.GetAddressOf());

    // get device context
    WRL::ComPtr<ID3D11DeviceContext> deviceContext;
    device->GetImmediateContext(deviceContext.GetAddressOf());
        
    uint32 levelWidth = textureDesc.Width >> mipLevel;
    if (levelWidth == 0)
    {
        levelWidth = 1;
    }

    uint32 levelHeight = textureDesc.Height >> mipLevel;
    if (levelHeight == 0)
    {
        levelHeight = 1;
    }

    win32_result wr;
        
    // data validation
    if ((destX + dataWidth) > levelWidth ||
        (destY + dataHeight) > levelHeight)
    {
        wr = E_INVALIDARG;
    }

    WRL::ComPtr<ID3D11Texture2D> textureToLock;
    uint32 levelToLock = 0;
    bool canWriteLock = false;
    bool needToComputeDestOffsets = false;

    // Can copy all the data if the destination texture size matches the data size, and we 
    // are referencing all the incoming data (not just a part of it).
        
    bool copyAllData = 
        (dataX == 0 && dataY == 0 && dataWidth == levelWidth && dataHeight == levelHeight) &&
        ((pitch * dataHeight) == dataSizeInBytes);

    canWriteLock = ((textureDesc.CPUAccessFlags & D3D11_CPU_ACCESS_WRITE) == D3D11_CPU_ACCESS_WRITE);
    if (canWriteLock)
    {
        textureToLock = _texture;
        levelToLock = mipLevel;
        needToComputeDestOffsets = true;

        // In addition, we need to make sure we are copying to the whole texture since we are 
        // working off the real (live) texture and not a surrogate.
        copyAllData = copyAllData && (destX == 0 && destY == 0);
    }
    else
    {
        // Only going to create a surrogate texture for the level we're locking, 
        // so the texture we end up copying from just has 1 mip level.
                
        levelToLock = 0;
                
        // Don't need to compute offsets because we create the surrogate texture 
        // as the exact dimensions of the data we are writing.
                
        needToComputeDestOffsets = false;

        D3D11_TEXTURE2D_DESC surrogateDesc;
        surrogateDesc.Width = dataWidth;
        surrogateDesc.Height = dataHeight;
        surrogateDesc.MipLevels = 1;
        surrogateDesc.ArraySize = 1;
        surrogateDesc.Format = textureDesc.Format;
        surrogateDesc.SampleDesc.Count = textureDesc.SampleDesc.Count;
        surrogateDesc.SampleDesc.Quality = textureDesc.SampleDesc.Quality;
        surrogateDesc.Usage = D3D11_USAGE_STAGING;
        surrogateDesc.BindFlags = 0;
        surrogateDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ|D3D11_CPU_ACCESS_WRITE;
        surrogateDesc.MiscFlags = 0;

        wr = device->CreateTexture2D(&surrogateDesc, NULL, textureToLock.GetAddressOf());
    }
        
    bool isLocked = false;
    D3D11_MAPPED_SUBRESOURCE mappedResource;
    
    withFinally::run(
    [&]()
    {
        wr = deviceContext->Map(
            textureToLock.Get(),
            levelToLock,
            D3D11_MAP_WRITE,
            0,
            &mappedResource
            );
            
        isLocked = true;
                
        if (mappedResource.RowPitch > pitch)
        {
            // D3D11 padded the texture, so we need to copy row by row 
            // to only copy the actual data.
            copyAllData = false;
        }

        if (copyAllData)
        {
            memcpy(mappedResource.pData, data, dataSizeInBytes);
        }
        else
        {
            uint32 numBytesToCopyPerRow = dataBytesPerPixel * dataWidth;
            for (uint32 j = 0; j < dataHeight; j++)
            {
                uint32 destIndex = 0;
                if (needToComputeDestOffsets)
                {
                    uint32 destBytesPerPixel = mappedResource.RowPitch / levelWidth;
                    destIndex = ((destY + j) * mappedResource.RowPitch) + (destX * destBytesPerPixel);
                }
                else
                {
                    destIndex = j * mappedResource.RowPitch;
                }

                uint32 srcIndex = ((dataY + j) * pitch) + (dataX * dataBytesPerPixel);
                memcpy((memory*)mappedResource.pData + destIndex, data + srcIndex, numBytesToCopyPerRow);
            }
        }
    },
    [&]()
    {
        if (isLocked)
        {
            deviceContext->Unmap(textureToLock.Get(), levelToLock);
        }
    });

    if (!canWriteLock)
    {
        deviceContext->CopySubresourceRegion(
            _texture.Get(),
            mipLevel,
            dataX,
            dataY,
            0,
            textureToLock.Get(),
            levelToLock,
            NULL
            );
    }
}
    
// gets bits per pixel
inline uint32 texture2d::getBitsPerPixel()
{
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);

    return ::baja::graphics::dx11_1::details::getBitsPerPixel(desc.Format);
}
    
// locks texture
inline void texture2d::lock( 
    const uint32 level,
    const lock_type lockType,
    memory** data,
    uint32* pitch
    ) 
{
    BAJA_CHECK_ARG(data != nullptr);
    BAJA_CHECK_ARG(pitch != nullptr);

    win32_result wr;

    // get texture description
    D3D11_TEXTURE2D_DESC textureDesc;
    _texture->GetDesc(&textureDesc);

    // get device
    WRL::ComPtr<ID3D11Device> device;
    _texture->GetDevice(device.GetAddressOf());

    // get device context
    WRL::ComPtr<ID3D11DeviceContext> deviceContext;
    device->GetImmediateContext(deviceContext.GetAddressOf());

    // can't lock if already locked
    if (_locked) wr = E_FAIL; // ERROR: 

    // can't lock a level that doesn't exist
    if (level >= textureDesc.MipLevels) wr = E_FAIL; // ERROR:

    // assume lock will not use surrogate data
    _lockUsingSurrogate = false;

    if (lockType == lock_type::read)
    {
        // are we able to take a read lock?
        bool canReadLock = ((textureDesc.CPUAccessFlags & D3D11_CPU_ACCESS_READ) == D3D11_CPU_ACCESS_READ);
        if (!canReadLock) wr = E_FAIL; // ERROR:
    }
    else
    {
        // are we able to take a write lock directly or need to use a surrogate?
        bool canWriteLock = ((textureDesc.CPUAccessFlags & D3D11_CPU_ACCESS_WRITE) == D3D11_CPU_ACCESS_WRITE);
        _lockUsingSurrogate = !canWriteLock;
    }

    WRL::ComPtr<ID3D11Texture2D> textureToLock;
    D3D11_TEXTURE2D_DESC textureToLockDesc = textureDesc;
    uint32 textureToLockLevel = level;

    if (!_lockUsingSurrogate)
    {
        // we can lock against the real texture
        textureToLock = _texture;
    }
    else
    {
        uint32 levelWidth = textureDesc.Width >> level;
        if (levelWidth == 0)
        {
            levelWidth = 1;
        }

        uint32 levelHeight = textureDesc.Height >> level;
        if (levelHeight == 0)
        {
            levelHeight = 1;
        }

        // ensure that no surrogate already exists
        if (!!_surrogate) wr = E_FAIL; // ERROR:

        // create surrogate surface used for locking
        D3D11_TEXTURE2D_DESC surrogateDesc;
        surrogateDesc.Width = levelWidth;
        surrogateDesc.Height = levelHeight;
        surrogateDesc.MipLevels = 1;
        surrogateDesc.ArraySize = 1;
        surrogateDesc.Format = textureDesc.Format;
        surrogateDesc.SampleDesc.Count = textureDesc.SampleDesc.Count;
        surrogateDesc.SampleDesc.Quality = textureDesc.SampleDesc.Quality;
        surrogateDesc.Usage = D3D11_USAGE_STAGING;
        surrogateDesc.BindFlags = 0;
        surrogateDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ|D3D11_CPU_ACCESS_WRITE;
        surrogateDesc.MiscFlags = 0;

        wr = device->CreateTexture2D(&surrogateDesc, NULL, _surrogate.GetAddressOf());
        textureToLock = _surrogate;
        textureToLockDesc = surrogateDesc;
        textureToLockLevel = 0;
    }

    // we should have a valid texture to lock 
    if (textureToLock.Get() == NULL) wr = E_POINTER;

    D3D11_MAP mappingType;
    if (lockType == lock_type::read)
    {
        mappingType = D3D11_MAP_READ;
    }
    else if (textureToLockDesc.Usage == D3D11_USAGE_DYNAMIC)
    {
        mappingType = D3D11_MAP_WRITE_DISCARD;
    }
    else
    {
        mappingType = D3D11_MAP_WRITE;
    }

    // take the lock
    D3D11_MAPPED_SUBRESOURCE mappedResource;
    wr = deviceContext->Map(
        textureToLock.Get(),
        textureToLockLevel,
        mappingType,
        0,
        &mappedResource);

    *data = (memory*)mappedResource.pData;
    *pitch = mappedResource.RowPitch;

    // remember level that was locked
    _lockedLevel = level;
    _locked = true;
}
        
// unlocks texture
inline void texture2d::unlock()
{
    win32_result wr;

    // get texture description
    D3D11_TEXTURE2D_DESC desc;
    _texture->GetDesc(&desc);

    // get device
    WRL::ComPtr<ID3D11Device> device;
    _texture->GetDevice(device.GetAddressOf());

    // get device context
    WRL::ComPtr<ID3D11DeviceContext> deviceContext;
    device->GetImmediateContext(deviceContext.GetAddressOf());

    if (!_locked) wr = E_FAIL; // ERROR:

    if (!_lockUsingSurrogate)
    {
        deviceContext->Unmap(_texture.Get(), _lockedLevel);
    }
    else
    {
        deviceContext->Unmap(_surrogate.Get(), 0);
        deviceContext->CopySubresourceRegion(
            _texture.Get(),
            _lockedLevel,
            0,
            0,
            0,
            _surrogate.Get(),
            0,
            NULL);
    }

    _surrogate.Reset();
    _locked = false;
    _lockedLevel = 0;
    _lockUsingSurrogate = false;
}

inline void texture2d::initialize(ID3D11Texture2D* texture2d)
{
    BAJA_CHECK_ARG(texture2d != nullptr);

    _views.clear();
    _texture = texture2d;
}
    
inline void texture2d::initialize(ID3D11View* view)
{
    BAJA_CHECK_ARG(view != nullptr);

    _views.clear();
    _views.push_back(view);

    WRL::ComPtr<ID3D11Resource> resource;
    _views[0]->GetResource(resource.GetAddressOf());

    win32_result((HRESULT)resource.As(&_texture));
}
    
inline void texture2d::initialize(const std::vector<WRL::ComPtr<ID3D11View>>& views)
{
    BAJA_CHECK_ARG(views.size() > 0);

    _views = views;

    WRL::ComPtr<ID3D11Resource> resource;

    for (size_t i = 0; i < views.size(); i++)
    {
        if (_views[i] != nullptr)
        {
            _views[i]->GetResource(&resource);
            break;
        }
    }

    win32_result((HRESULT)resource.As(&_texture));
}

// Creates an ID3D11Texture2D object
inline WRL::ComPtr<ID3D11Texture2D> createTexture2DObject(
    ID3D11Device* d3d11Device,
    const uint32 width,
    const uint32 height,
    const uint32 mipLevels,
    const pixel_format format,
    const uint32 msaaCount,
    const uint32 msaaQualityLevel,
    const D3D11_USAGE usage,
    const uint32 binding,
    const uint32 cpuAccess,
    const uint32 miscFlags
    )
{ 
    BAJA_CHECK_ARG(d3d11Device != NULL);

    // initialize the texture description
    D3D11_TEXTURE2D_DESC desc;
    desc.Width = width;
    desc.Height = height;
    desc.MipLevels = mipLevels;
    desc.ArraySize = 1;
    desc.Format = toDXGI_FORMAT(format);
    desc.SampleDesc.Count = msaaCount;
    desc.SampleDesc.Quality = msaaQualityLevel;
    desc.Usage = usage;
    desc.BindFlags = binding;
    desc.CPUAccessFlags = cpuAccess;
    desc.MiscFlags = miscFlags;

    WRL::ComPtr<ID3D11Texture2D> texture;
    // create new texture
    win32_result((HRESULT)d3d11Device->CreateTexture2D(&desc, NULL, texture.GetAddressOf()));

    return texture;
}

}}}} // namespace baja::graphics::dx11_1::details
